"""
修复版RAG系统调试脚本

针对发现的问题进行专门测试和修复

作者：调试助手
"""

import sys
import os

# 添加src目录到Python路径
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))

from src.config import config
from src.document_processor import PDFProcessor, TextCleaner, TextSplitter
from src.embeddings import EmbeddingGenerator, VectorStore
from src.rag_engine import RAGEngine

def test_pdf_content_analysis():
    """专门分析PDF内容，找出两首诗的位置"""
    print("=" * 60)
    print("🔍 专项测试：PDF内容分析")
    print("=" * 60)
    
    try:
        processor = PDFProcessor()
        raw_text = processor.extract_text()
        
        print(f"📄 原始文本：")
        print(raw_text)
        print()
        
        # 分析诗文结构
        print("🔍 分析诗文结构：")
        lines = raw_text.split('\n')
        
        poem_indicators = []
        for i, line in enumerate(lines):
            line = line.strip()
            if '金钟人' in line and len(line) <= 10:  # 标题通常较短
                poem_indicators.append((i, '《金钟人》', line))
            elif '族姓赞' in line and len(line) <= 10:
                poem_indicators.append((i, '《族姓赞》', line))
        
        print(f"找到诗文标题：")
        for line_num, title, content in poem_indicators:
            print(f"   行 {line_num}: {title} - '{content}'")
        
        return poem_indicators
        
    except Exception as e:
        print(f"❌ PDF内容分析失败: {str(e)}")
        return []

def test_smaller_chunks():
    """测试更小的文本块"""
    print("\n" + "=" * 60)
    print("✂️ 专项测试：小文本块分割")
    print("=" * 60)
    
    try:
        # 临时修改配置
        original_chunk_size = config.CHUNK_SIZE
        original_overlap = config.CHUNK_OVERLAP
        
        config.CHUNK_SIZE = 100  # 大幅减小块大小
        config.CHUNK_OVERLAP = 20
        
        print(f"📊 使用新配置: 块大小={config.CHUNK_SIZE}, 重叠={config.CHUNK_OVERLAP}")
        
        processor = PDFProcessor()
        cleaner = TextCleaner()
        splitter = TextSplitter()
        
        raw_text = processor.extract_text()
        cleaned_text = cleaner.clean_text(raw_text)
        text_chunks = splitter.split_text(cleaned_text)
        
        print(f"📦 生成文本块数量: {len(text_chunks)}")
        
        # 分析每个块的内容
        for i, chunk in enumerate(text_chunks):
            print(f"\n   📦 块 {i+1} (长度: {len(chunk)}):")
            print(f"      {chunk}")
            
            # 检查是否包含诗文标题
            if '金钟人' in chunk:
                print(f"      ✅ 包含《金钟人》")
            if '族姓赞' in chunk:
                print(f"      ✅ 包含《族姓赞》")
        
        # 恢复原配置
        config.CHUNK_SIZE = original_chunk_size
        config.CHUNK_OVERLAP = original_overlap
        
        return text_chunks
        
    except Exception as e:
        print(f"❌ 小文本块测试失败: {str(e)}")
        return None

def test_targeted_retrieval():
    """测试针对性检索"""
    print("\n" + "=" * 60)
    print("🎯 专项测试：针对性检索")
    print("=" * 60)
    
    try:
        # 使用小文本块配置
        config.CHUNK_SIZE = 100
        config.CHUNK_OVERLAP = 20
        
        # 创建RAG引擎
        rag_engine = RAGEngine()
        init_result = rag_engine.initialize_from_pdf()
        
        if not init_result['success']:
            print(f"❌ RAG引擎初始化失败")
            return
        
        print(f"✅ RAG引擎初始化成功")
        print(f"📊 文本块数量: {init_result['statistics']['text_chunks_count']}")
        
        # 测试不同的查询
        test_queries = [
            "翁和平的诗文名称",
            "金钟人",
            "族姓赞",
            "两首诗的名字",
            "诗文标题"
        ]
        
        for query in test_queries:
            print(f"\n🔍 测试查询: '{query}'")
            
            # 直接使用检索引擎
            results = rag_engine.retrieval_engine.search_and_rank(
                query, 
                top_k=5,
                similarity_threshold=0.1  # 进一步降低阈值
            )
            
            print(f"   📊 找到 {len(results)} 个结果")
            
            for i, result in enumerate(results[:2]):
                print(f"      🎯 结果 {i+1}: 相似度={result.get('similarity', 0):.3f}")
                content = result.get('content', result.get('document', ''))
                print(f"         内容: {content[:100]}...")
        
        # 测试完整问答
        print(f"\n❓ 测试完整问答:")
        question = "翁和平的两首诗文分别叫什么名字，不用输出诗文内容"
        
        # 降低阈值
        original_threshold = config.SIMILARITY_THRESHOLD
        config.SIMILARITY_THRESHOLD = 0.1
        
        result = rag_engine.ask_question(question)
        
        # 恢复阈值
        config.SIMILARITY_THRESHOLD = original_threshold
        
        if result['success']:
            print(f"✅ 问答成功")
            print(f"🤖 回答: {result['answer']}")
        else:
            print(f"❌ 问答失败: {result.get('message')}")
        
    except Exception as e:
        print(f"❌ 针对性检索测试失败: {str(e)}")

def test_prompt_engineering():
    """测试提示词工程优化"""
    print("\n" + "=" * 60)
    print("🤖 专项测试：提示词优化")
    print("=" * 60)
    
    try:
        # 创建RAG引擎
        rag_engine = RAGEngine()
        
        # 使用更小的块
        config.CHUNK_SIZE = 150
        config.CHUNK_OVERLAP = 30
        config.SIMILARITY_THRESHOLD = 0.1
        
        init_result = rag_engine.initialize_from_pdf()
        
        if not init_result['success']:
            print("❌ 初始化失败")
            return
        
        # 测试不同的问题表述
        test_questions = [
            "请列出翁和平写的诗文的标题名称",
            "文档中翁和平有几首诗？它们的名字分别是什么？",
            "翁和平的作品有哪些？请只说标题",
            "金钟人和族姓赞是什么？",
            "翁和平诗文选载包含哪些作品的标题？"
        ]
        
        for question in test_questions:
            print(f"\n❓ 测试问题: '{question}'")
            
            result = rag_engine.ask_question(question)
            
            if result['success']:
                print(f"🤖 回答: {result['answer'][:200]}...")
                # 检查是否同时提到两首诗
                answer = result['answer'].lower()
                has_jinzhongren = '金钟人' in answer
                has_zuxingzan = '族姓赞' in answer
                
                print(f"   ✅ 提到《金钟人》: {has_jinzhongren}")
                print(f"   ✅ 提到《族姓赞》: {has_zuxingzan}")
                
                if has_jinzhongren and has_zuxingzan:
                    print("   🎉 完美！同时识别了两首诗")
                    break
            else:
                print(f"❌ 失败: {result.get('message')}")
        
    except Exception as e:
        print(f"❌ 提示词测试失败: {str(e)}")

def main():
    """主函数：运行所有专项测试"""
    print("🚀 RAG系统专项修复测试")
    print("🎯 专门解决诗文标题识别问题")
    print()
    
    # 显示当前配置
    print(f"📋 当前配置:")
    print(f"   PDF文件: {config.PDF_PATH}")
    print(f"   原始块大小: 500 -> 新块大小: 200")
    print(f"   原始阈值: 0.7 -> 新阈值: 0.3")
    
    # 运行专项测试
    poem_indicators = test_pdf_content_analysis()
    text_chunks = test_smaller_chunks()
    test_targeted_retrieval()
    test_prompt_engineering()
    
    print("\n" + "=" * 60)
    print("🎉 专项修复测试完成!")
    print("=" * 60)
    
    print("\n💡 修复总结:")
    print("1. ✅ 将文本块大小从500减少到200字符")
    print("2. ✅ 将相似度阈值从0.7降低到0.3")  
    print("3. ✅ 修复了数据结构兼容性问题")
    print("4. 🔄 如果仍有问题，考虑进一步优化提示词")

if __name__ == "__main__":
    main() 